package com.inspur.edp.ai.flow.node;

import com.inspur.edp.ai.flow.exception.FlowException;
import com.inspur.edp.ai.flow.metadata.AbsNode;
import com.inspur.edp.ai.flow.metadata.Parameter;
import com.inspur.edp.ai.flow.spi.NodeCompiler;
import com.inspur.edp.ai.flow.spi.GetOutputParameter;
import com.inspur.edp.common.expr.Express;
import com.inspur.edp.common.expr.ExpressDSLGeneratorMgr;
import com.inspur.edp.common.expr.ast.ASTContext;
import com.inspur.edp.common.expr.ast.CompileUnitParser;
import com.inspur.edp.common.expr.evaluator.EvaluatorService;
import com.inspur.edp.common.expr.evaluator.ExpressEvaluator;
import com.inspur.edp.common.expr.parser.ASTParser;
import com.inspur.edp.common.expr.parser.ExpressAstParser;
import com.inspur.edp.common.type.TypeRefer;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author lizhaorui
 * @date 2025/9/22
 * @description
 */
public class LoopNodeCompiler implements NodeCompiler, GetOutputParameter {

    @Override
    public String getKind() {
        return LoopNode.KIND;
    }

    @Override
    public Map<String, ExpressEvaluator> compile(ASTContext context, AbsNode absNode) {
        LoopNode loopNode = (LoopNode) absNode;
        loopNode.setInputSelf(true);
        Map<String, ExpressEvaluator> needInputMap = new HashMap<>();
        Express iterableExpr = loopNode.getIterableExpr();
        if (iterableExpr != null) {
            ExpressEvaluator expressEvaluator = EvaluatorService.build(context, iterableExpr);
            loopNode.setIterableExprEvaluator(expressEvaluator);
            //处理输入参数的依赖
            EvaluatorService.putDependVars(context, iterableExpr, needInputMap);
        }
        return needInputMap;
    }

    @Override
    public List<Parameter> getOutputParams(ASTContext astContext, AbsNode node) {
        List<Parameter> parameters = new ArrayList<>();
        if (node.getOutputParams() != null) {
            parameters.addAll(node.getOutputParams());
        }
        LoopNode loopNode = (LoopNode) node;
        Express iterableExpr = loopNode.getIterableExpr();
        String dsl = ExpressDSLGeneratorMgr.generateDsl(astContext, iterableExpr);
        CompileUnitParser.ExpressContext expressContext = ASTParser.parse(dsl).express();
        TypeRefer iterableExprType = ExpressAstParser.getOriginalType(expressContext, astContext);
        // 添加默认的 item 参数
        Parameter itemParam = new Parameter();
        itemParam.setCode(loopNode.getIterableVariable());
        List<TypeRefer> genericTypes = iterableExprType.getGenericTypes();
        if (genericTypes == null || genericTypes.size() == 0) {
            throw new FlowException("iterable express must is iterable type.");
        }
        itemParam.setType(genericTypes.get(0));
        parameters.add(itemParam);
        // 添加默认的 index 参数
        Parameter indexParam = new Parameter();
        indexParam.setCode(LoopNode.ITEM_INDEX_KEY);
        indexParam.setType(TypeRefer.INT_TYPE);
        parameters.add(indexParam);

        return parameters;
    }
}